{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "overview",
   "metadata": {},
   "source": [
    "# Strands Agent with AgentCore Memory Tutorial using Hooks\n",
    "\n",
    "## Overview\n",
    "\n",
    "This tutorial demonstrates how to build an intelligent personal assistant using Strands agents integrated with AgentCore Memory through hooks. The agent maintains conversation context and learns from interactions to provide personalized responses.\n",
    "\n",
    "## Tutorial Details\n",
    "\n",
    "**Use Case**: Math Assistant\n",
    "\n",
    "| Information         | Details                                                                          |\n",
    "|:--------------------|:---------------------------------------------------------------------------------|\n",
    "| Tutorial type       | Long term Conversational                                                         |\n",
    "| Agent type          | Math Assistant                                                                   |\n",
    "| Agentic Framework   | Strands Agents                                                                   |\n",
    "| LLM model           | Anthropic Claude Sonnet 3.7                                                      |\n",
    "| Tutorial components | AgentCore Summary Strategy for Memory, Hooks for storing and retrieving Memory   |\n",
    "| Example complexity  | Intermediate                                                                     |\n",
    "\n",
    "\n",
    "You'll learn to:\n",
    "- Set up AgentCore Memory with conversation summaries\n",
    "- Create memory hooks for automatic storage and retrieval\n",
    "- Build a Strands agent with persistent memory\n",
    "- Test memory functionality across conversations\n",
    "\n",
    "### Scenario Context\n",
    "\n",
    "In this example you'll create a Math Assistant example where you'd store summaries of the previous conversations. \n",
    "Key features of this example:\n",
    "- **Automatic Memory Storage**: Conversations are automatically saved\n",
    "- **Context Retrieval**: Previous conversations inform current responses\n",
    "- **Summary Generation**: Key information is extracted and summarized\n",
    "- **Tool Integration**: Calculator tool for mathematical operations\n",
    "\n",
    "## Architecture\n",
    "<div style=\"text-align:left\">\n",
    "    <img src=\"architecture.png\" width=\"65%\" />\n",
    "</div>\n",
    "\n",
    "\n",
    "## Prerequisites\n",
    "\n",
    "To execute this tutorial you will need:\n",
    "- Python 3.10+\n",
    "- AWS credentials with Amazon Bedrock AgentCore Memory permissions\n",
    "- Amazon Bedrock AgentCore SDK"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "917e6ed5-ca98-4a5a-a569-3a1ab68998e5",
   "metadata": {},
   "source": [
    "## Step 1: Environment set up\n",
    "Let's begin importing all the necessary libraries and defining the clients to make this notebook work."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "404db75f-614e-4152-af9a-5b28afb56e33",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install -qr requirements.txt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0a2e1911-a9fa-4578-868b-a22bff8ab298",
   "metadata": {},
   "outputs": [],
   "source": [
    "from bedrock_agentcore.memory import MemoryClient\n",
    "from bedrock_agentcore.memory.constants import StrategyType"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "imports",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import logging\n",
    "from strands import Agent\n",
    "from datetime import datetime\n",
    "from strands_tools import calculator\n",
    "from strands.hooks import AfterInvocationEvent, HookProvider, HookRegistry, MessageAddedEvent\n",
    "\n",
    "# Setup logging\n",
    "logging.basicConfig(level=logging.INFO, format=\"%(asctime)s - %(levelname)s - %(message)s\")\n",
    "logger = logging.getLogger(\"memory-tutorial\")\n",
    "\n",
    "# Configuration - replace with your values\n",
    "REGION = os.getenv('AWS_REGION', 'us-west-2')\n",
    "ROLE_ARN = \"<<INSERT-YOUR-IAM-ROLE>>\"\n",
    "ACTOR_ID = f\"actor-{datetime.now().strftime('%Y%m%d%H%M%S')}\"\n",
    "SESSION_ID = f\"tutorial-{datetime.now().strftime('%Y%m%d%H%M%S')}\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "step2-desc",
   "metadata": {},
   "source": [
    "## Step 2: Create Memory Resource\n",
    "\n",
    "In this step, we're creating our memory resource with a summary strategy. This resource will store and organize our conversation data. The strategy we're defining will automatically generate summaries of conversations and store them in organized namespaces.\n",
    "\n",
    "Firstly, lets create a custom prompt for the math assistant.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a5315e29",
   "metadata": {},
   "outputs": [],
   "source": [
    "CUSTOM_PROMPT = \"\"\"\n",
    "Your task is to extract math learning data from the user's conversations. You store the progress of the user in a memory system to understand their math level and help them progress.\n",
    "\n",
    "You are tasked with analyzing conversations to extract the user's math learning patterns. You'll be analyzing two sets of data: \n",
    "\n",
    "<past_conversation> \n",
    "[Past conversations between the user and math tutor will be placed here for context] \n",
    "</past_conversation> \n",
    "\n",
    "<current_conversation> \n",
    "[The current conversation between the user and math tutor will be placed here] \n",
    "</current_conversation> \n",
    "\n",
    "Your job is to identify and categorize the user's math learning profile:\n",
    "- Extract the user's current math level from problems they solve correctly/incorrectly\n",
    "- Extract the user's preferred learning style from how they ask questions and respond to explanations\n",
    "- Extract topic strengths and weaknesses from their performance patterns\n",
    "- Track learning progress and identify areas needing reinforcement\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "create_memory",
   "metadata": {},
   "outputs": [],
   "source": [
    "from botocore.exceptions import ClientError\n",
    "\n",
    "# Initialize Memory Client\n",
    "client = MemoryClient(region_name=REGION)\n",
    "memory_name = \"MathAssistant\"\n",
    "# Define memory strategy for conversation summaries\n",
    "strategies = [\n",
    "    {\n",
    "        StrategyType.CUSTOM.value: {\n",
    "            \"name\": \"CustomSemanticMemory\",\n",
    "            \"description\": \"Captures facts from conversations\",\n",
    "            \"namespaces\": [\"/students/math/{actorId}\"],\n",
    "            \"configuration\" : {\n",
    "                \"semanticOverride\" : {\n",
    "                    \"extraction\" : {\n",
    "                        \"modelId\" : \"anthropic.claude-3-5-sonnet-20241022-v2:0\",\n",
    "                        \"appendToPrompt\": CUSTOM_PROMPT\n",
    "                    }\n",
    "                },\n",
    "    }}}\n",
    "]\n",
    "\n",
    "# Create memory resource\n",
    "try:\n",
    "    memory = client.create_memory_and_wait(\n",
    "        name=memory_name,\n",
    "        strategies=strategies, # Use the defined long term strategies\n",
    "        description=\"Memory for tutorial agent\",\n",
    "        event_expiry_days=30,\n",
    "        memory_execution_role_arn=ROLE_ARN,\n",
    "    )\n",
    "    memory_id = memory['id']\n",
    "    logger.info(f\"✅ Created memory: {memory_id}\")\n",
    "except ClientError as e:\n",
    "    if e.response['Error']['Code'] == 'ValidationException' and \"already exists\" in str(e):\n",
    "        # If memory already exists, retrieve its ID\n",
    "        memories = client.list_memories()\n",
    "        memory_id = next((m['id'] for m in memories if m['id'].startswith(memory_name)), None)\n",
    "        logger.info(f\"Memory already exists. Using existing memory ID: {memory_id}\")\n",
    "except Exception as e:\n",
    "    # Handle any errors during memory creation\n",
    "    logger.info(f\"❌ ERROR: {e}\")\n",
    "    import traceback\n",
    "    traceback.print_exc()\n",
    "    # Cleanup on error - delete the memory if it was partially created\n",
    "    if memory_id:\n",
    "        try:\n",
    "            client.delete_memory_and_wait(memory_id=memory_id)\n",
    "            logger.info(f\"Cleaned up memory: {memory_id}\")\n",
    "        except Exception as cleanup_error:\n",
    "            logger.info(f\"Failed to clean up memory: {cleanup_error}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "step3-desc",
   "metadata": {},
   "source": [
    "## Step 3: Create Memory Hook Provider\n",
    "\n",
    "This step defines our custom `MemoryHookProvider` class that automates memory operations. Hooks are special functions that run at specific points in an agent's execution lifecycle. The memory hook we're creating serves two primary functions:\n",
    "\n",
    "1. **Retrieve Memories**: Automatically fetches relevant past conversations when a user sends a message\n",
    "2. **Save Memories**: Stores new conversations after the agent responds\n",
    "\n",
    "This creates a seamless memory experience without manual management."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "memory_hooks",
   "metadata": {},
   "outputs": [],
   "source": [
    "class MemoryHookProvider(HookProvider):\n",
    "    \"\"\"Hook provider for automatic memory management\"\"\"\n",
    "    \n",
    "    def __init__(self, memory_id: str, client: MemoryClient, actor_id: str, session_id: str):\n",
    "        self.memory_id = memory_id\n",
    "        self.client = client\n",
    "        self.actor_id = actor_id\n",
    "        self.session_id = session_id\n",
    "        self.namespace = f\"/students/math/{self.actor_id}\"\n",
    "    \n",
    "    def retrieve_memories(self, event: MessageAddedEvent):\n",
    "        \"\"\"Retrieve relevant memories before processing user message\"\"\"\n",
    "        messages = event.agent.messages\n",
    "        if messages[-1][\"role\"] == \"user\" and \"toolResult\" not in messages[-1][\"content\"][0]:\n",
    "            user_message = messages[-1][\"content\"][0].get(\"text\", \"\")\n",
    "            \n",
    "            try:\n",
    "                # Retrieve relevant memories\n",
    "                memories = self.client.retrieve_memories(\n",
    "                    memory_id=self.memory_id,\n",
    "                    namespace=self.namespace,\n",
    "                    query=user_message\n",
    "                )\n",
    "                \n",
    "                # Extract memory content\n",
    "                memory_context = []\n",
    "                for memory in memories:\n",
    "                    if isinstance(memory, dict):\n",
    "                        content = memory.get('content', {})\n",
    "                        if isinstance(content, dict):\n",
    "                            text = content.get('text', '').strip()\n",
    "                            if text:\n",
    "                                memory_context.append(text)\n",
    "                \n",
    "                # Inject memories into user message\n",
    "                if memory_context:\n",
    "                    context_text = \"\\n\".join(memory_context)\n",
    "                    original_text = messages[-1][\"content\"][0].get(\"text\", \"\")\n",
    "                    messages[-1][\"content\"][0][\"text\"] = (\n",
    "                        f\"{original_text}\\n\\nPrevious context: {context_text}\"\n",
    "                    )\n",
    "                    logger.info(f\"Retrieved {len(memory_context)} memories\")\n",
    "                    \n",
    "            except Exception as e:\n",
    "                logger.error(f\"Failed to retrieve memories: {e}\")\n",
    "    \n",
    "    def save_memories(self, event: AfterInvocationEvent):\n",
    "        \"\"\"Save conversation after agent response\"\"\"\n",
    "        try:\n",
    "            messages = event.agent.messages\n",
    "            if len(messages) >= 2 and messages[-1][\"role\"] == \"assistant\":\n",
    "                # Get last user and assistant messages\n",
    "                user_msg = None\n",
    "                assistant_msg = None\n",
    "                \n",
    "                for msg in reversed(messages):\n",
    "                    if msg[\"role\"] == \"assistant\" and not assistant_msg:\n",
    "                        assistant_msg = msg[\"content\"][0][\"text\"]\n",
    "                    elif msg[\"role\"] == \"user\" and not user_msg and \"toolResult\" not in msg[\"content\"][0]:\n",
    "                        user_msg = msg[\"content\"][0][\"text\"]\n",
    "                        break\n",
    "                \n",
    "                if user_msg and assistant_msg:\n",
    "                    # Save conversation\n",
    "                    self.client.create_event(\n",
    "                        memory_id=self.memory_id,\n",
    "                        actor_id=self.actor_id,\n",
    "                        session_id=self.session_id,\n",
    "                        messages=[(user_msg, \"USER\"), (assistant_msg, \"ASSISTANT\")]\n",
    "                    )\n",
    "                    logger.info(\"Saved conversation to memory\")\n",
    "                    \n",
    "        except Exception as e:\n",
    "            logger.error(f\"Failed to save memories: {e}\")\n",
    "    \n",
    "    def register_hooks(self, registry: HookRegistry) -> None:\n",
    "        \"\"\"Register memory hooks\"\"\"\n",
    "        registry.add_callback(MessageAddedEvent, self.retrieve_memories)\n",
    "        registry.add_callback(AfterInvocationEvent, self.save_memories)\n",
    "        logger.info(\"Memory hooks registered\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "step4-desc",
   "metadata": {},
   "source": [
    "## Step 4: Create Agent with Memory\n",
    "\n",
    "Now we're creating our Strands agent and connecting it with our memory hook provider. This agent will have two key capabilities:\n",
    "\n",
    "1. **Memory Integration**: The memory hooks we created will enable automatic context retrieval\n",
    "2. **Calculator Tool**: The agent can perform mathematical operations when needed\n",
    "\n",
    "This combination creates a personal assistant that both remembers past interactions and can perform useful calculations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "create_agent",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create memory hook provider\n",
    "memory_hooks = MemoryHookProvider(\n",
    "    memory_id=memory_id,\n",
    "    client=client,\n",
    "    actor_id=ACTOR_ID,\n",
    "    session_id=SESSION_ID\n",
    ")\n",
    "\n",
    "# Create agent with memory hooks and calculator tool\n",
    "agent = Agent(\n",
    "    hooks=[memory_hooks],\n",
    "    tools=[calculator],\n",
    "    system_prompt=\"You are a helpful personal math tutor. You assist users in solving math problems and provide personalized assistance.\"\n",
    ")\n",
    "\n",
    "print(\"✅ Agent created with memory hooks.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "step5-desc",
   "metadata": {},
   "source": [
    "**We have our agent set up ! Let's test it now.**\n",
    "\n",
    "## Test Memory Functionality\n",
    "\n",
    "In this section, we'll test the agent's memory capabilities through a series of interactions. We'll observe how the agent builds context over time and recalls previous interactions.\n",
    "\n",
    "First, let's introduce ourselves to the agent and ask a math question:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "test1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# First interaction - introduce yourself\n",
    "response1 = agent(\"Hi, I'm John and I just enrolled in Discrete Math course. Help me solve this: How many ways can I arrange 5 books on a shelf?\")\n",
    "print(f\"Agent: {response1}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "test3-desc",
   "metadata": {},
   "source": [
    "Let's give the agent another calculation task:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "test3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Second interaction - another calculation\n",
    "response2 = agent(\"I learn better with step-by-step explanation with example questions. Can you explain modular arithmetic? What's 17 mod 5?\")\n",
    "print(f\"Agent: {response2}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "test2-desc",
   "metadata": {},
   "source": [
    "Now, let's see if the agent remembers who we are.\n",
    "\n",
    "**Note:** Give a ~20 sec pause here to allow some time for the memory to be extracted, consolidated and stored."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "test2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Third interaction - test memory recall\n",
    "response3 = agent(\"I got that right! What's the immediate next step that I should study after modular arithmetic?\")\n",
    "print(f\"Agent: {response3}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "test4-desc",
   "metadata": {},
   "source": [
    "Finally, let's check if the agent remembers our calculation history:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "test4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Fourth interaction - test context awareness\n",
    "response4 = agent(\"This is too hard, can we try something easier?\")\n",
    "print(f\"Agent: {response4}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "step6-desc",
   "metadata": {},
   "source": [
    "### Verify Memory Storage\n",
    "\n",
    "As a final step, we'll verify that our conversations have been properly stored in AgentCore Memory. This demonstrates that the memory hooks are working correctly and the agent can access this information in future interactions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "verify_memory",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Check stored memories\n",
    "try:\n",
    "    memories = client.retrieve_memories(\n",
    "        memory_id=memory_id,\n",
    "        namespace=f\"/students/math/{ACTOR_ID}\",\n",
    "        query=\"mathematics calculations\"\n",
    "    )\n",
    "    \n",
    "    print(f\"\\n📚 Found {len(memories)} memories:\")\n",
    "    for i, memory in enumerate(memories, 1):\n",
    "        if isinstance(memory, dict):\n",
    "            content = memory.get('content', {})\n",
    "            if isinstance(content, dict):\n",
    "                text = content.get('text', '')[:200] + \"...\"\n",
    "                print(f\"{i}. {text}\")\n",
    "                \n",
    "except Exception as e:\n",
    "    print(f\"Error retrieving memories: {e}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d0822184",
   "metadata": {},
   "source": [
    "Tutorial completed! 🎉\n",
    "\n",
    "Key takeaways:\n",
    "- Memory hooks automatically store and retrieve conversation context\n",
    "- Agents can maintain state across multiple interactions\n",
    "- AgentCore Memory provides semantic search for relevant context\n",
    "- Tools can be combined with memory for enhanced functionality"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cleanup-desc",
   "metadata": {},
   "source": [
    "## Clean Up\n",
    "\n",
    "### Optional: Delete Memory Resource\n",
    "\n",
    "After completing the tutorial, you may want to delete the memory resource to avoid incurring unnecessary costs. The following code is provided for cleanup but is commented out by default."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cleanup_code",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Uncomment to delete the memory resource\n",
    "# try:\n",
    "#     client.delete_memory_and_wait(memory_id=memory_id)\n",
    "#     print(f\"✅ Deleted memory resource: {memory_id}\")\n",
    "# except Exception as e:\n",
    "#     print(f\"Error deleting memory: {e}\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv2",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.13.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
